# 该文件包含webui通用工具，可以被不同的webui使用
from typing import *
from pathlib import Path
from typing import Iterator

from httpx import Response

from configs.client_config import (

    SCORE_THRESHOLD,
    VECTOR_SEARCH_TOP_K,
)
import httpx
import contextlib
import json
from io import BytesIO



def set_httpx_timeout(timeout=60.0):
    '''
    设置httpx默认timeout到60秒。
    httpx默认timeout是5秒，在请求LLM回答时不够用。
    '''
    httpx._config.DEFAULT_TIMEOUT_CONFIG.connect = timeout
    httpx._config.DEFAULT_TIMEOUT_CONFIG.read = timeout
    httpx._config.DEFAULT_TIMEOUT_CONFIG.write = timeout


# KB_ROOT_PATH = Path(KB_ROOT_PATH)
set_httpx_timeout()


class ApiRequest:
    '''
    api.py调用的封装,主要实现:
    1. 简化api调用方式
    2. 实现无api调用(直接运行server.chat.*中的视图函数获取结果),无需启动api.py
    '''

    def __init__(
            self,
            base_url: str = "http://127.0.0.1:7861",
            timeout: float = 60.0,
            no_remote_api: bool = False,  # True to call api view function directly
    ):
        self.base_url = base_url
        self.timeout = timeout
        self.no_remote_api = no_remote_api

    def _parse_url(self, url: str) -> str:
        if (not url.startswith("http")
                and self.base_url
        ):
            part1 = self.base_url.strip(" /")
            part2 = url.strip(" /")
            return f"{part1}/{part2}"
        else:
            return url

    def get(
            self,
            url: str,
            params: Union[Dict, List[Tuple], bytes] = None,
            retry: int = 3,
            stream: bool = False,
            **kwargs: Any,
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        while retry > 0:
            try:
                if stream:
                    return httpx.stream("GET", url, params=params, **kwargs)
                else:
                    return httpx.get(url, params=params, **kwargs)
            except Exception as e:
                print('Exception:',e)
                retry -= 1

    async def aget(
            self,
            url: str,
            params: Union[Dict, List[Tuple], bytes] = None,
            retry: int = 3,
            stream: bool = False,
            **kwargs: Any,
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        async with httpx.AsyncClient() as client:
            while retry > 0:
                try:
                    if stream:
                        return await client.stream("GET", url, params=params, **kwargs)
                    else:
                        return await client.get(url, params=params, **kwargs)
                except Exception as e:
                    print('Exception:',e)
                    retry -= 1

    def post(
            self,
            url: str,
            data: Dict = None,
            json: Dict = None,
            retry: int = 3,
            stream: bool = False,
            **kwargs: Any
    ) -> Union[httpx.Response,Iterator[Response], None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        while retry > 0:
            try:
                # return requests.post(url, data=data, json=json, stream=stream, **kwargs)
                if stream:
                    return httpx.stream("POST", url, data=data, json=json, **kwargs)
                else:
                    return httpx.post(url, data=data, json=json, **kwargs)
            except Exception as e:
                print('Exception:',e)
                retry -= 1

    async def apost(
            self,
            url: str,
            data: Dict = None,
            json: Dict = None,
            retry: int = 3,
            stream: bool = False,
            **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        async with httpx.AsyncClient() as client:
            while retry > 0:
                try:
                    if stream:
                        return await client.stream("POST", url, data=data, json=json, **kwargs)
                    else:
                        return await client.post(url, data=data, json=json, **kwargs)
                except Exception as e:
                    print('Exception:',e)
                    retry -= 1

    def delete(
            self,
            url: str,
            data: Dict = None,
            json: Dict = None,
            retry: int = 3,
            stream: bool = False,
            **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        while retry > 0:
            try:
                if stream:
                    return httpx.stream("DELETE", url, data=data, json=json, **kwargs)
                else:
                    return httpx.delete(url, data=data, json=json, **kwargs)
            except Exception as e:
                print('Exception:',e)
                retry -= 1

    async def adelete(
            self,
            url: str,
            data: Dict = None,
            json: Dict = None,
            retry: int = 3,
            stream: bool = False,
            **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        async with httpx.AsyncClient() as client:
            while retry > 0:
                try:
                    if stream:
                        return await client.stream("DELETE", url, data=data, json=json, **kwargs)
                    else:
                        return await client.delete(url, data=data, json=json, **kwargs)
                except Exception as e:
                    print('Exception:',e)
                    retry -= 1


    def _httpx_stream2generator(
            self,
            response: contextlib._GeneratorContextManager,
            as_json: bool = False,
    ):
        '''
        将httpx.stream返回的GeneratorContextManager转化为普通生成器
        '''
        try:
            with response as r:
                for chunk in r.iter_text(None):
                    if as_json and chunk:
                        yield json.loads(chunk)
                    elif chunk.strip():
                        yield chunk
        except httpx.ConnectError as e:
            msg = f"无法连接API服务器，请确认 ‘api.py’ 已正常启动。"
            print('Exception:',msg)
            print('Exception:',e)
            yield {"code": 500, "msg": msg}
        except httpx.ReadTimeout as e:
            msg = f"API通信超时，请确认已启动FastChat与API服务（详见RADME）"
            print('Exception:',msg)
            print('Exception:',e)
            yield {"code": 500, "msg": msg}
        except Exception as e:
            print('Exception:',e)
            yield {"code": 500, "msg": str(e)}

    def chat_chat(
            self,
            query: str,
            history: List[Dict] = [],
            stream: bool = True,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/chat/chat接口
        '''
        data = {
            "query": query,
            "history": history,
            "stream": stream,
        }

        response = self.post("/chat/chat", json=data, stream=True)
        return self._httpx_stream2generator(response)

    def knowledge_base_chat(
            self,
            query: str,
            knowledge_base_name: str,
            top_k: int = VECTOR_SEARCH_TOP_K,
            score_threshold: float = SCORE_THRESHOLD,
            history: List[Dict] = [],
            stream: bool = True,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/chat/knowledge_base_chat接口
        '''

        data = {
            "query": query,
            "knowledge_base_name": knowledge_base_name,
            "top_k": top_k,
            "score_threshold": score_threshold,
            "history": history,
            "stream": stream,
            "local_doc_url": no_remote_api,
        }

        response = self.post(
            "/chat/knowledge_base_chat",
            json=data,
            stream=True,
        )
        return self._httpx_stream2generator(response, as_json=True)


    # 知识库相关操作

    def _check_httpx_json_response(
            self,
            response: httpx.Response,
            errorMsg: str = f"无法连接API服务器，请确认已执行python server\\api.py",
    ) -> Dict:
        '''
        check whether httpx returns correct data with normal Response.
        error in api with streaming support was checked in _httpx_stream2enerator
        '''
        try:
            return response.json()
        except Exception as e:
            print('Exception:',e)
            return {"code": 500, "msg": errorMsg or str(e)}

    def list_knowledge_bases(
            self,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/list_knowledge_bases接口
        '''

        response = self.get("/knowledge_base/list_knowledge_bases")
        data = self._check_httpx_json_response(response)
        return data.get("data", [])

    def create_knowledge_base(
            self,
            knowledge_base_name: str,
            embed_model: str,
            vector_store_type: str = "faiss",
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/create_knowledge_base接口
        '''

        data = {
            "knowledge_base_name": knowledge_base_name,
            "vector_store_type": vector_store_type,
            "embed_model": embed_model,
        }

        response = self.post(
            "/knowledge_base/create_knowledge_base",
            json=data,
        )

        return self._check_httpx_json_response(response)

    def delete_knowledge_base(
            self,
            knowledge_base_name: str,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/delete_knowledge_base接口
        '''

        response = self.post(
            "/knowledge_base/delete_knowledge_base",
            json=f"{knowledge_base_name}",
        )
        return self._check_httpx_json_response(response)

    def list_kb_docs(
            self,
            knowledge_base_name: str,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/list_docs接口
        '''

        response = self.get(
            "/knowledge_base/list_docs",
            params={"knowledge_base_name": knowledge_base_name}
        )
        data = self._check_httpx_json_response(response)
        return data.get("data", [])

    def file_exists(
            self,
            kb: str,
            selected_rows: list,
            no_remote_api: bool = None,
    ) :
        '''
        对应api.py/knowledge_base/file_exists接口
        '''

        response = self.post(
            "/knowledge_base/file_exists",
            json={
                "kb": kb,
                "selected_rows": selected_rows
            }
        )
        if response.headers['content-type']== 'multipart/form-data':
            # 是后端传回的FileResponse(media_type="multipart/form-data")
            # 获取文件名和内容

            def extract_filename_from_header(header:str):
                import re
                from urllib.parse import unquote
                match = re.search(r'filename\*?=(?:UTF-8\'\'|utf-8\'\')?([^;\n]+)', header)

                if match:
                    match_result =match.group(1)
                    match_result = unquote(match_result)  # 解码URL编码
                    return match_result
                else:
                    return None

            file_name = extract_filename_from_header(response.headers["content-disposition"])
            if not file_name:
                file_name = '后端找不到本文件'
            file_content = response.content

        else:
            file_name=''
            file_content=None
        return file_name,file_content

    def upload_kb_doc(
            self,
            file: Union[str, Path, bytes],
            knowledge_base_name: str,
            filename: str = None,
            override: bool = False,
            not_refresh_vs_cache: bool = False,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/upload_doc接口
        '''

        if isinstance(file, bytes):  # raw bytes
            file = BytesIO(file)
        elif hasattr(file, "read"):  # a file io like object
            filename = filename or file.name
        else:  # a local path
            file = Path(file).absolute().open("rb")
            filename = filename or file.name

        response = self.post(
            "/knowledge_base/upload_doc",
            data={
                "knowledge_base_name": knowledge_base_name,
                "override": override,
                "not_refresh_vs_cache": not_refresh_vs_cache,
            },
            files={"file": (filename, file)},
        )
        return self._check_httpx_json_response(response)

    def delete_kb_doc(
            self,
            knowledge_base_name: str,
            doc_name: str,
            delete_content: bool = False,
            not_refresh_vs_cache: bool = False,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/delete_doc接口
        '''

        data = {
            "knowledge_base_name": knowledge_base_name,
            "doc_name": doc_name,
            "delete_content": delete_content,
            "not_refresh_vs_cache": not_refresh_vs_cache,
        }

        response = self.post(
            "/knowledge_base/delete_doc",
            json=data,
        )
        return self._check_httpx_json_response(response)

    def update_kb_doc(
            self,
            knowledge_base_name: str,
            file_name: str,
            not_refresh_vs_cache: bool = False,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/update_doc接口
        '''

        response = self.post(
            "/knowledge_base/update_doc",
            json={
                "knowledge_base_name": knowledge_base_name,
                "file_name": file_name,
                "not_refresh_vs_cache": not_refresh_vs_cache,
            },
        )
        return self._check_httpx_json_response(response)

    def recreate_vector_store(
            self,
            knowledge_base_name: str,
            vs_type: str ,
            embed_model: str,
            allow_empty_kb: bool = True,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/recreate_vector_store接口
        '''

        data = {
            "knowledge_base_name": knowledge_base_name,
            "allow_empty_kb": allow_empty_kb,
            "vs_type": vs_type,
            "embed_model": embed_model,
        }

        response = self.post(
            "/knowledge_base/recreate_vector_store",
            json=data,
            stream=True,
            timeout=None,
        )
        return self._httpx_stream2generator(response, as_json=True)

    def get_kb_details(
            self,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/get_kb_details接口
        '''

        response = self.get("/knowledge_base/get_kb_details")
        data = self._check_httpx_json_response(response)
        return data.get("data", [])

    def get_kb_doc_details(
            self,
            knowledge_base_name: str,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/get_kb_doc_details接口
        '''

        response = self.get(
            "/knowledge_base/get_kb_doc_details",
            params={"knowledge_base_name": knowledge_base_name}
        )
        data = self._check_httpx_json_response(response)
        return data.get("data", [])

    # file_agent相关
    def upload_files_for_agent(
            self,
            file: Union[str, Path, bytes],
            filename: str = None,
            override: bool = False,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/file_agent/upload_file接口
        '''
        if isinstance(file, bytes):  # raw bytes
            file = BytesIO(file)
        elif hasattr(file, "read"):  # a file io like object
            filename = filename or file.name
        else:  # a local path
            file = Path(file).absolute().open("rb")
            filename = filename or file.name

        response = self.post(
            "/file_agent/upload_file",
            data={
                "override": override,
            },
            files={"file": (filename, file)},
        )
        return self._check_httpx_json_response(response)

    def create_agent_for_files(
            self,
    ):
        '''
        对应api.py/file_agent/create_agent接口
        '''

        response = self.get(
            "/file_agent/create_agent",
        )
        return self._check_httpx_json_response(response)

    def ask_file_agent(
            self,
            query: str,
            stream: bool = True,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/file_agent/ask_file_agent接口
        '''
        data = {
            "query": query,
            "stream": stream,
        }

        response = self.post("/file_agent/ask_file_agent", json=data, stream=stream)
        if isinstance(response, httpx.Response) or isinstance(response,contextlib._GeneratorContextManager):
            if stream:

                return self._httpx_stream2generator(response)
            else:
                data= self._check_httpx_json_response(response)
                return data.get("data", [''])[0]
        else:
            return None

    def list_docs_in_file_agent(
            self,
            no_remote_api: bool = None,
    ):
        '''
        对应api.py/file_agent/list_docs_in_file_agent接口
        '''

        response = self.get("/file_agent/list_docs_in_file_agent")
        data = self._check_httpx_json_response(response)
        return data.get("data", [])

    def delete_docs_in_file_agent(
            self,
    ):
        '''
        对应api.py/file_agent/delete_docs_in_file_agent接口
        '''

        response = self.get(
            "/file_agent/delete_docs_in_file_agent",
        )
        return self._check_httpx_json_response(response)

    #启动前端时要从后端加载的东西
    def get_server_config(
            self,
    ):
        '''
        对应api.py/get_server_config接口
        '''

        response = self.get("/get_server_config")
        data = self._check_httpx_json_response(response)
        return data.get("data", [])





def check_error_msg(data: Union[str, dict, list], key: str = "errorMsg") -> str:
    '''
    return error message if error occured when requests API
    '''
    if isinstance(data, dict):
        if key in data:
            return data[key]
        if "code" in data and data["code"] != 200:
            return data["msg"]
    return ""


def check_success_msg(data: Union[str, dict, list], key: str = "msg") -> str:
    '''
    return error message if error occured when requests API
    '''
    if (isinstance(data, dict)
            and key in data
            and "code" in data
            and data["code"] == 200):
        return data[key]
    return ""


if __name__ == "__main__":
    api = ApiRequest(no_remote_api=True)
